package com.cl.spark.node;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.cl.spark.base.BaseSparkNode;
import com.cl.spark.dto.SparkParam;
import com.cl.spark.dto.SparkResult;
import com.cl.spark.enums.SparkNodeEnum;

import org.apache.spark.sql.*;
import org.apache.spark.sql.types.DataTypes;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import scala.collection.JavaConverters;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@Component
public class FilterNode extends BaseSparkNode {
    @Autowired
    SparkSession sparkSession;

    @Override
    public SparkResult process(SparkParam sparkParam) {
        Dataset<Row> dataset = sparkParam.getSparkResultList().get(0).getDataset();
        JSONObject expression = sparkParam.getNodeExpression();

        JSONArray fields = expression.getJSONArray("fields");
        JSONArray fields2 = expression.getJSONArray("fields2");
        String filterFunction = expression.getString("filterFunction");
        String filterValue = expression.getString("filterValue");
        String logicalOperator = expression.getString("logicalOperator"); // 默认 AND

        List<Column> filterConditions = new ArrayList<>();
        for (int i = 0; i < fields.size(); i++) {
            String field = fields.getString(i);
            if (fields2 != null && !fields2.isEmpty() && i < fields2.size()) {
                String field2 = fields2.getString(i);
                filterConditions.add(generateFilterColumn2(field, filterFunction, field2));
            } else {
                filterConditions.add(generateFilterColumn(field, filterFunction, filterValue));
            }
        }

        Column combinedCondition;
        if ("OR".equalsIgnoreCase(logicalOperator)) {
            combinedCondition = filterConditions.stream()
                    .reduce((c1, c2) -> c1.or(c2))
                    .orElse(functions.lit(false)); // 如果没有条件，则返回一个总是为假的条件
        } else {
            combinedCondition = filterConditions.stream()
                    .reduce((c1, c2) -> c1.and(c2))
                    .orElse(functions.lit(true)); // 如果没有条件，则返回一个总是为真的条件
        }
        Dataset<Row> filter = dataset.filter(combinedCondition);
        return SparkResult.success(filter);

    }

    public Column generateFilterColumn(String field, String filterFunction, String filterValue) {
        if (!StringUtils.hasLength(field) || !StringUtils.hasLength(filterFunction)) {
            throw new RuntimeException(String.format("field: %s function %s value %s", field, filterFunction, filterValue));
        }
        Column column = new Column(field);
        Column filterValueLit = functions.lit(filterValue);
        switch (filterFunction) {
            case "大于":
                column = column.gt(filterValueLit);
                break;
            case "小于":
                column = column.lt(filterValueLit);
                break;
            case "不包含":
                column = functions.not(column.contains(filterValue));
                break;
            case "包含":
                column = column.contains( filterValue);
                break;
            case "等于":
                column = column.equalTo(filterValueLit);
                break;
            case "不等于":
                column = column.notEqual(filterValueLit);
                break;
            case "大于等于":
                column = column.geq(filterValueLit);
                break;
            case "小于等于":
                column = column.leq(filterValueLit);
                break;
            case "列表包含(或)":
                column = column.rlike(filterValue.replace(",", "|"));
                break;
            case "空值":
                column = column.isNull().or(functions.length(column).equalTo(0));
                break;
            case "非空值":
                column = column.isNotNull().and(functions.length(column).gt(0));
                break;
            case "非纯数字":
                column = column.isNaN();
                break;
            case "以值为开头":
                column = column.startsWith(filterValueLit);
                break;
            case "以值为结尾":
                column = column.endsWith(filterValueLit);
                break;
        }
        return column;
    }

    public Column generateFilterColumn2(String field, String filterFunction, String field2) {
        if (!StringUtils.hasLength(field) || !StringUtils.hasLength(filterFunction)) {
            throw new RuntimeException(String.format("field: %s function %s value %s", field, filterFunction, field2));
        }
        Column column = new Column(field);
        Column filterValueLit = new Column(field2);
        switch (filterFunction) {
            case "大于":
                column = column.gt(filterValueLit);
                break;
            case "小于":
                column = column.lt(filterValueLit);
                break;
            case "不包含":
                column = functions.not(column.contains(filterValueLit));
                break;
            case "包含":
                column = column.contains(filterValueLit);
                break;
            case "等于":
                column = column.equalTo(filterValueLit);
                break;
            case "不等于":
                column = column.notEqual(filterValueLit);
                break;
            case "大于等于":
                column = column.geq(filterValueLit);
                break;
            case "小于等于":
                column = column.leq(filterValueLit);
                break;
            case "空值":
                column = column.isNull().or(functions.length(column).equalTo(0));
                break;
            case "非空值":
                column = column.isNotNull().and(functions.length(column).gt(0));
                break;
            case "非纯数字":
                column = column.isNaN();
                break;
            case "以值为开头":
                column = column.startsWith(filterValueLit);
                break;
            case "以值为结尾":
                column = column.endsWith(filterValueLit);
                break;
        }
        return column;
    }

    @Override
    public SparkNodeEnum getType() {
        return SparkNodeEnum.FILTER;
    }
}
